import RPi.GPIO as GPIO
import time

# 各轮速度补偿系数（0.0~2.0）
coeff_lf = 1.0  # 左前轮系数
coeff_rf = 1.0 # 右前轮系数
coeff_lb = 1.0  # 左后轮系数
coeff_rb = 1.0 # 右后轮系数
# 定义L298N的引脚（已避开18和23号引脚）
IN1_LF = 27
IN2_LF = 17
ENA_LF = 22
IN1_RF = 24
IN2_RF = 25
ENA_RF = 12
IN1_LB = 6
IN2_LB = 5
ENA_LB = 13
IN1_RB = 19
IN2_RB = 16
ENA_RB = 26 

# 初始化PWM对象
pwm_lf = None
pwm_rf = None
pwm_lb = None
pwm_rb = None


def setup():
    # 设置GPIO模式为BCM
    GPIO.setmode(GPIO.BCM)

    # 设置引脚模式
    GPIO.setup(IN1_LF, GPIO.OUT)
    GPIO.setup(IN2_LF, GPIO.OUT)
    GPIO.setup(ENA_LF, GPIO.OUT)

    GPIO.setup(IN1_RF, GPIO.OUT)
    GPIO.setup(IN2_RF, GPIO.OUT)
    GPIO.setup(ENA_RF, GPIO.OUT)

    GPIO.setup(IN1_LB, GPIO.OUT)
    GPIO.setup(IN2_LB, GPIO.OUT)
    GPIO.setup(ENA_LB, GPIO.OUT)

    GPIO.setup(IN1_RB, GPIO.OUT)
    GPIO.setup(IN2_RB, GPIO.OUT)
    GPIO.setup(ENA_RB, GPIO.OUT)

    # 初始化PWM（频率设置为1000Hz）
    global pwm_lf, pwm_rf, pwm_lb, pwm_rb
    pwm_lf = GPIO.PWM(ENA_LF, 1000)
    pwm_rf = GPIO.PWM(ENA_RF, 1000)
    pwm_lb = GPIO.PWM(ENA_LB, 1000)
    pwm_rb = GPIO.PWM(ENA_RB, 1000)

    # 启动PWM
    pwm_lf.start(0)
    pwm_rf.start(0)
    pwm_lb.start(0)
    pwm_rb.start(0)
def setMotorSpeed(in1, in2, pwm, speed):
    duty_cycle = abs(speed) * 100 / 255
    if speed > 0:
        GPIO.output(in1, GPIO.HIGH)
        GPIO.output(in2, GPIO.LOW)
    elif speed < 0:
        GPIO.output(in1, GPIO.LOW)
        GPIO.output(in2, GPIO.HIGH)
    else:
        GPIO.output(in1, GPIO.LOW)
        GPIO.output(in2, GPIO.LOW)
    pwm.ChangeDutyCycle(duty_cycle)
def setWheelSpeeds(Vx, Vy, Vr):
    # 计算理论速度
    Vlf = Vy + Vx + Vr
    Vrf = Vy - Vx - Vr
    Vlb = Vy - Vx + Vr
    Vrb = Vy + Vx - Vr

    # 应用各轮补偿系数
    Vlf = int(Vlf * coeff_lf)
    Vrf = int(Vrf * coeff_rf)
    Vlb = int(Vlb * coeff_lb)
    Vrb = int(Vrb * coeff_rb)

    # 归一化处理
    max_speed = max(abs(Vlf), abs(Vrf), abs(Vlb), abs(Vrb))
    if max_speed > 255:
        scale = 255 / max_speed
        Vlf = int(Vlf * scale)
        Vrf = int(Vrf * scale)
        Vlb = int(Vlb * scale)
        Vrb = int(Vrb * scale)

    # 设置电机速度
    setMotorSpeed(IN1_LF, IN2_LF, pwm_lf, Vlf)
    setMotorSpeed(IN1_RF, IN2_RF, pwm_rf, Vrf)
    setMotorSpeed(IN1_LB, IN2_LB, pwm_lb, Vlb)
    setMotorSpeed(IN1_RB, IN2_RB, pwm_rb, Vrb)
def stopRobot():
    setMotorSpeed(IN1_LF, IN2_LF, pwm_lf, 0)
    setMotorSpeed(IN1_RF, IN2_RF, pwm_rf, 0)
    setMotorSpeed(IN1_LB, IN2_LB, pwm_lb, 0)
    setMotorSpeed(IN1_RB, IN2_RB, pwm_rb, 0)
def 前进(duration, speed):
    setWheelSpeeds(0, speed, 0)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 后退(duration, speed):
    setWheelSpeeds(0, -speed, 0)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 向左移动(duration, speed):
    setWheelSpeeds(-speed, 0, 0)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 向右移动(duration, speed):
    setWheelSpeeds(speed, 0, 0)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 向左前45度移动(duration, speed):
    setWheelSpeeds(-speed, speed, 0)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 向右后45度移动(duration, speed):
    setWheelSpeeds(speed, -speed, 0)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 左转90度(duration, speed):
    # 设置纯旋转（Vx=0, Vy=0, Vr=正速度）
    setWheelSpeeds(Vx=0, Vy=0, Vr=speed)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 右转90度(duration, speed):
    # 设置纯旋转（Vx=0, Vy=0, Vr=负速度）
    setWheelSpeeds(Vx=0, Vy=0, Vr=-speed)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)
def 原地转180度(duration, speed):
    # 使用与左转相同的逻辑，但通常需要更长的duration
    setWheelSpeeds(Vx=0, Vy=0, Vr=speed)
    time.sleep(duration / 1000.0)
    stopRobot()
    time.sleep(1)